{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ef331be9",
   "metadata": {
    "id": "ef331be9"
   },
   "source": [
    "![nebullvm nebuly AI accelerate inference optimize DeepLearning](https://user-images.githubusercontent.com/38586138/201391643-a80407e5-2c28-409c-90c9-327795cd27e8.png)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f260653a",
   "metadata": {
    "id": "f260653a"
   },
   "source": [
    "# Accelerate Stable Diffusion with Speedster\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "8bdf3af5",
   "metadata": {
    "id": "8bdf3af5"
   },
   "source": [
    "Hi and welcome 👋\n",
    "\n",
    "In this notebook we will discover how in just a few steps you can speed up the response time of Stable Diffusion inference using the Speedster module from the open-source library nebullvm. In the first section we will try using `Speedster` with the default configuration, then we will explore a more advanced option that involves the TensorRT plugins, that allow to accelerate Stable Diffusion further on GPU.\n",
    "\n",
    "Let's jump to the code."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cXXh1ifQ13mH",
   "metadata": {
    "id": "cXXh1ifQ13mH"
   },
   "source": [
    "# Installation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48aljCHu14-H",
   "metadata": {
    "id": "48aljCHu14-H"
   },
   "source": [
    "Install Speedster:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "QFQh3BVr1-GO",
   "metadata": {
    "id": "QFQh3BVr1-GO"
   },
   "outputs": [],
   "source": [
    "!pip install speedster"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a7a86b3",
   "metadata": {
    "id": "8a7a86b3"
   },
   "source": [
    "Install deep learning compilers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cffbfa32",
   "metadata": {
    "id": "cffbfa32"
   },
   "outputs": [],
   "source": [
    "!python -m nebullvm.installers.auto_installer --frameworks diffusers --compilers all"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "c2ab3de7",
   "metadata": {},
   "source": [
    "# Environment check (GPU only)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "61a1a445",
   "metadata": {},
   "source": [
    "**Please skip this section if you don't have a GPU**"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e2784bb8",
   "metadata": {},
   "source": [
    "If you want to optimize Stable Diffusion on a Nvidia GPU, in order to work properly, the following requirements must be installed on your machine:\n",
    "- `CUDA>=12.0`\n",
    "- `tensorrt>=8.6.0`\n",
    "- `torch<=1.13.1`"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e3bc8b4d",
   "metadata": {},
   "source": [
    "From TensorRT 8.6, all the tensorrt pre-built wheels released by nvidia support only `CUDA>=12.0`. Speedster will install `tensorrt>=8.6.0` automatically in the auto-installer only if it detects CUDA>=12.0, otherwise it will install `tensorrt==8.5.3.1`. In that case, you will have to upgrade your CUDA version and then to upgarde tensorrt to 8.6.0 or above to execute this notebook.\n",
    "\n",
    "There should be a way to run TensorRT 8.6 also with CUDA 11, but it requires installing TensorRT in a different way, you can check this issue: https://github.com/NVIDIA/TensorRT/issues/2773. Otherwise, we highly suggest to just upgrade to CUDA 12.\n",
    "\n",
    "For now PyTorch>=2.0.0 is not supported due to an [issue](https://github.com/pytorch/pytorch/issues/97262) in the conversion to onnx, so until they fix it you must have torch<=1.13.1 to optimize Stable Diffusion successfully."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ec2267f0",
   "metadata": {},
   "source": [
    "First of all, Let's check the CUDA version installed on the machine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82b78585",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import subprocess\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    cuda_version = subprocess.check_output([\"nvidia-smi\"])\n",
    "    cuda_version = int(cuda_version.decode(\"utf-8\").split(\"\\n\")[2].split(\"|\")[-2].split(\":\")[-1].strip().split(\".\")[0])\n",
    "    assert cuda_version >= 12, (\"This notebook requires CUDA>=12.0 to be executed, please upgrade your CUDA version.\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "015cfa92",
   "metadata": {},
   "source": [
    "If you have CUDA<12.0, you can upgrade it at this link: https://developer.nvidia.com/cuda-downloads"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "563779e6",
   "metadata": {},
   "source": [
    "Then, let's check the tensorrt version installed on the platform. Stable Diffusion optimization is supported starting from `tensorrt==8.6.0`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e385021d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorrt\n",
    "from nebullvm.tools.utils import check_module_version\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    assert check_module_version(tensorrt, \"8.6.0\"), (\"This notebook can be run only with tensorrt>=8.6.0, if using an older version you could have issues during the optimization. Please upgrade your version.\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "61da505b",
   "metadata": {},
   "source": [
    "If you have an older version, after ensuring you have `CUDA>=12.0` installed, you can upgrade your TensorRT version by running:\n",
    "```\n",
    "pip install -U tensorrt\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3876bea4",
   "metadata": {},
   "source": [
    "Finally, let's check the PyTorch version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db83853f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from nebullvm.tools.utils import check_module_version\n",
    "\n",
    "assert check_module_version(torch, max_version=\"1.13.1+cu117\"), (\"This notebook can be run only with torch<=1.13.1, if using an older version you could have issues during the optimization. Please downgrade your version.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73072506",
   "metadata": {
    "id": "73072506"
   },
   "source": [
    "## Model and Dataset setup"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "aeb2c521",
   "metadata": {},
   "source": [
    "Once we have ensured that the the required libraries are installed, we have to choose the version of Stable Diffusion we want to optimize, speedster officially supports the most used versions:\n",
    "- `CompVis/stable-diffusion-v1-4`\n",
    "- `runwayml/stable-diffusion-v1-5`\n",
    "- `stabilityai/stable-diffusion-2-1-base`\n",
    "- `stabilityai/stable-diffusion-2-1` (only on gpus with at least 22GB of Memory, if you want to try with a GPU with a lower memory, you have to uncomment `pipe.enable_attention_slicing()` in the cell below)\n",
    "\n",
    "Other Stable Diffusion versions from the Diffusers library should work but have never been tested. If you try a version not included among these and it works, please feel free to report it to us on [Discord](https://discord.com/invite/RbeQMu886J) so we can add it to the list of supported versions. If you try a version that does not work, you can open an issue and possibly a PR on [GitHub](https://github.com/nebuly-ai/nebullvm/issues)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e4d55115",
   "metadata": {
    "id": "e4d55115"
   },
   "source": [
    "For this notebook, we are going to select Stable Diffusion 1.4. Let's download and load it using the diffusers API:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d633cf21",
   "metadata": {
    "id": "d633cf21",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from diffusers import StableDiffusionPipeline\n",
    "\n",
    "# Select Stable Diffusion version\n",
    "model_id = \"CompVis/stable-diffusion-v1-4\"\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "if device == \"cuda\":\n",
    "    # On GPU we load by default the model in half precision, because it's faster and lighter.\n",
    "    pipe = StableDiffusionPipeline.from_pretrained(model_id, revision='fp16', torch_dtype=torch.float16)\n",
    "    # pipe.enable_attention_slicing() # Uncomment for stable-diffusion-2.1 on gpus with 16GB of memory like V100-16GB and T4\n",
    "else:\n",
    "    pipe = StableDiffusionPipeline.from_pretrained(model_id)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "11aa0739",
   "metadata": {
    "id": "11aa0739"
   },
   "source": [
    "Let's now create an example dataset with some random sentences, that will be used later for the optimization process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbbfeeb2",
   "metadata": {
    "id": "cbbfeeb2"
   },
   "outputs": [],
   "source": [
    "input_data = [\n",
    "    \"a photo of an astronaut riding a horse on mars\",\n",
    "    \"a monkey eating a banana in a forest\",\n",
    "    \"white car on a road surrounded by palm trees\",\n",
    "    \"a fridge full of bottles of beer\",\n",
    "    \"madara uchiha throwing asteroids against people\"\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17040431",
   "metadata": {
    "id": "17040431"
   },
   "source": [
    "## Speed up inference with Speedster"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44ddc21d",
   "metadata": {
    "id": "44ddc21d"
   },
   "source": [
    "It's now time of improving a bit the performance in terms of speed. Let's use `Speedster`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9d934f6",
   "metadata": {
    "id": "f9d934f6"
   },
   "outputs": [],
   "source": [
    "from speedster import optimize_model, save_model, load_model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2799e3e3",
   "metadata": {},
   "source": [
    "Let's move the pipe back to CPU to save up GPU memory, `Speedster` will automatically move it back to GPU when required."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45220cf0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "\n",
    "# Move the pipe back to cpu\n",
    "pipe.to(\"cpu\")\n",
    "\n",
    "# Clean memory\n",
    "torch.cuda.empty_cache()\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76248033",
   "metadata": {
    "id": "76248033"
   },
   "source": [
    "Using Speedster is very simple and straightforward! Just use the `optimize_model` function and provide as input the model, some input data as example and the optimization time mode. Optionally a dynamic_info dictionary can be also provided, in order to support inputs with dynamic shape."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "75b339c3",
   "metadata": {},
   "source": [
    "**Optimisation of stable diffusion requires a lot of RAM. If you are running this notebook on google colab, make sure to use the high RAM option, otherwise the kernel may crash. If the kernel crashes also when using the high RAM option, please try adding also `\"torchscript\"` to the `ignore_compilers` list. \n",
    "If running on GPU, the optimization requires at least 16GB og GPU memory to exploit the best techniques for optimizing the model, otherwise it may fail with a Memory Error**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "zPC_EDwEJIM0",
   "metadata": {
    "id": "zPC_EDwEJIM0"
   },
   "outputs": [],
   "source": [
    "optimized_model = optimize_model(\n",
    "    model=pipe,\n",
    "    input_data=input_data,\n",
    "    optimization_time=\"unconstrained\",\n",
    "    ignore_compilers=[\"torch_tensor_rt\", \"tvm\"],  # Some compilers have issues with Stable Diffusion, so it's better to skip them.\n",
    "    metric_drop_ths=0.2,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "fdae59d2",
   "metadata": {},
   "source": [
    "If running on GPU, here you should obtain a speedup of about 124% on the UNet. We run the optimization on a **3090Ti** and here are our results:\n",
    "- **Original Model (PyTorch, fp16): 51,557 ms/batch**\n",
    "- **Optimized Model (TensorRT, fp16): 23,055 ms/batch**\n",
    "\n",
    "If the optimized model you obtained is not a TensorRT one, probably there was an error during the optimization. If running on colab, it could happen that the standard gpu is not enough to run the optimization, so we suggest to select a premium gpu with more memory.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "af9f86ac",
   "metadata": {},
   "source": [
    "If everything worked correctly, let's check the output of the optimized model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b640885",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_prompt = \"futuristic llama with a cyberpunk city on the background\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa443637",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_model(test_prompt).images[0]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6e5b3b21",
   "metadata": {
    "id": "6e5b3b21"
   },
   "source": [
    "Let's run the prediction 10 times to calculate the average response time of the original model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09170c78",
   "metadata": {},
   "outputs": [],
   "source": [
    "if device == \"cuda\":\n",
    "    pipe = StableDiffusionPipeline.from_pretrained(model_id, revision='fp16', torch_dtype=torch.float16)\n",
    "    # pipe.enable_attention_slicing() # Uncomment for stable-diffusion-2.1 on gpus with 16GB of memory like V100-16GB and T4\n",
    "else:\n",
    "    pipe = StableDiffusionPipeline.from_pretrained(model_id)\n",
    "\n",
    "pipe.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3bc5c98",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "d3bc5c98",
    "outputId": "e0596cf2-fa96-4c50-c012-f5cdab82e681"
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "times = []\n",
    "\n",
    "# Warmup for 2 iterations\n",
    "for _ in range(2):\n",
    "    with torch.no_grad():\n",
    "        final_out = pipe(test_prompt).images[0]\n",
    "\n",
    "# Benchmark\n",
    "for _ in range(8):\n",
    "    st = time.time()\n",
    "    with torch.no_grad():\n",
    "        final_out = pipe(test_prompt).images[0]\n",
    "    times.append(time.time()-st)\n",
    "original_model_time = sum(times)/len(times)\n",
    "print(f\"Average response time for original Stable Diffusion 1.4: {original_model_time} s\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3db0a7a1",
   "metadata": {
    "id": "3db0a7a1"
   },
   "source": [
    "Let's run the prediction 10 times to calculate the average response time of the optimized model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3e83997",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "a3e83997",
    "outputId": "7a416b14-f170-4df9-d416-026f06a7d980"
   },
   "outputs": [],
   "source": [
    "times = []\n",
    "\n",
    "for _ in range(2):\n",
    "    with torch.no_grad():\n",
    "        final_out = optimized_model(test_prompt).images[0]\n",
    "\n",
    "# Benchmark\n",
    "for _ in range(8):\n",
    "    st = time.time()\n",
    "    with torch.no_grad():\n",
    "        final_out = optimized_model(test_prompt).images[0]\n",
    "    times.append(time.time()-st)\n",
    "optimized_model_time = sum(times)/len(times)\n",
    "print(f\"Average response time for optimized Stable Diffusion 1.4: {optimized_model_time} s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ceb60d8c",
   "metadata": {
    "id": "ceb60d8c"
   },
   "source": [
    "## Save and reload the optimized model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9eda1a0",
   "metadata": {},
   "source": [
    "We can easily save to disk the optimized model with the following line:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62b6fcbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "save_model(optimized_model, \"model_save_path\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c968d51",
   "metadata": {},
   "source": [
    "We can then load again the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1340c49",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimized_model = load_model(\"model_save_path\", pipe=pipe)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb234e5e",
   "metadata": {
    "id": "cb234e5e"
   },
   "source": [
    "Great! Was it easy? How are the results? Do you have any comments?\n",
    "Share your optimization results and thoughts with <a href=\"https://discord.gg/RbeQMu886J\" target=\"_blank\"> our community on Discord</a>, where we chat about Speedster and AI acceleration.\n",
    "\n",
    "Note that the acceleration of Speedster depends very much on the hardware configuration and your AI model. Given the same input model, Speedster can accelerate it by 10 times on some machines and perform poorly on others.\n",
    "\n",
    "If you want to learn more about how Speedster works, look at other tutorials and performance benchmarks, check out the links below or write to us on Discord."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b77ff2ac",
   "metadata": {
    "id": "b77ff2ac"
   },
   "source": [
    "<center> \n",
    "    <a href=\"https://discord.com/invite/RbeQMu886J\" target=\"_blank\" style=\"text-decoration: none;\"> Join the community </a> |\n",
    "    <a href=\"https://nebuly.gitbook.io/nebuly/welcome/questions-and-contributions\" target=\"_blank\" style=\"text-decoration: none;\"> Contribute to the library </a>\n",
    "</center>\n",
    "\n",
    "<center> \n",
    "    <a href=\"https://github.com/nebuly-ai/nebullvm/tree/main/apps/accelerate/speedster#key-concepts\" target=\"_blank\" style=\"text-decoration: none;\"> How speedster works </a> •\n",
    "    <a href=\"https://github.com/nebuly-ai/nebullvm/tree/main/apps/accelerate/speedster#documentation\" target=\"_blank\" style=\"text-decoration: none;\"> Documentation </a> •\n",
    "    <a href=\"https://github.com/nebuly-ai/nebullvm/tree/main/apps/accelerate/speedster#quick-start\" target=\"_blank\" style=\"text-decoration: none;\"> Quick start </a> \n",
    "</center>"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "provenance": []
  },
  "gpuClass": "premium",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.15"
  },
  "vscode": {
   "interpreter": {
    "hash": "4ca44071b2152bc556aa4c839392f76fd4b80aa39d34257f2d304fa0d1d8b7d9"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
